#!/usr/bin/env python
# Create: 2018/10/31
__author__ = '749B'
"""异步扫描端口的模块"""

import socket
import selectors
import time


class ScanTask(object):
    def __init__(self, host, start, end):
        self.host = host
        self.start = start
        self.end = end

        self.port = start

    def __str__(self):
        return "%s:(%s-%s)" % (self.host, self.start, self.end)


class AsyncScanPort(object):
    def __init__(self, start=1, end=65535, timeout=1, interval=0.01, pool=100):
        self.sel = selectors.DefaultSelector()
        self.start = start
        self.end = end
        self.timeout = timeout
        self.interval = interval
        self.pool = pool

        self.scan_list = []
        self.scan_index = 0

    def add_scan_task(self, host, start=None, end=None):
        start = start or self.start
        end = end or self.end
        obj = ScanTask(host, start, end)
        self.scan_list.append(obj)

    def check_timeout(self):
        for i in self.sel.get_map().values():
            t = i.data['time']
            if time.time() > t + self.timeout:
                self.sel.unregister(i.fileobj)
                return True
            # 只要判断第一个就好了
            return False

    def add_to_scan(self):
        if len(self.sel.get_map()) >= self.pool:
            res = self.check_timeout()
            if not res:
                return False
        if self.scan_index >= len(self.scan_list):
            return False
        obj = self.scan_list[self.scan_index]
        if obj.port > obj.end:
            self.scan_index += 1
            print("\r单个地址加载完毕:", obj)
            return self.add_to_scan()
        client = socket.socket()
        client.setblocking(False)
        try:
            client.connect((obj.host, obj.port))
        except BlockingIOError as e:
            pass
        self.sel.register(client, selectors.EVENT_WRITE, {'obj': obj, 'port': obj.port, 'time': time.time()})
        print('\r正在扫描 ==> %s:%s' % (obj.host, obj.port), end='', flush=True)
        obj.port += 1
        return True

    def run(self):
        """主函数"""
        self.add_to_scan()
        while self.sel.get_map():
            events = self.sel.select(self.interval)
            if events:
                for key, mask in events:
                    self.callback(key.fileobj, mask, key.data)  # key.fileobj就是sel.register里第一次参数
            self.add_to_scan()
        print("扫描完毕")
        self.sel.close()

    def callback(self, sk, mask, data):
        try:
            sk.getpeername()
        except OSError as e:
            # 没有扫描到端口
            self.sel.unregister(sk)
            return
        self.sel.unregister(sk)
        print("\r扫描到端口:", data['obj'], data['port'])


if __name__ == '__main__':
    hosts = ['192.168.1.1', '192.168.1.2', '192.168.1.3']
    obj = AsyncScanPort()
    for h in hosts:
        obj.add_scan_task(h)
    obj.run()
